package com.gitee.xhs.saveAndModifyingPointWriter.broadcastState;

import com.gitee.xhs.EnvironmentConfig;
import com.gitee.xhs.FlinkEnvironmentFactory;
import com.gitee.xhs.broadcastState.BroadcastStateDao;
import com.gitee.xhs.dao.ModifyingPointDao;
import com.gitee.xhs.util.SavePointApiUtil;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.state.api.OperatorTransformation;
import org.apache.flink.state.api.SavepointWriter;
import org.apache.flink.state.api.StateBootstrapTransformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import static com.gitee.xhs.StateConstant.CHECK_POINT_BASE_PATH;
import static com.gitee.xhs.broadcastState.BroadcastStateConstant.BROADCAST_STATE_UID;
import static com.gitee.xhs.broadcastState.BroadcastStateConstant.CUSTOM_BROADCAST_STATE_DIR;
import static com.gitee.xhs.saveAndModifyingPointWriter.SavePointWriterConstant.MAX_PARALLELISM;
import static com.gitee.xhs.saveAndModifyingPointWriter.SavePointWriterConstant.MODIFYING_UID_PREFIX;

/**
 * @author BigHead
 * @date 2022/7/11 17:28
 * @description BroadcastStateNewSavePointsMain
 */
public class BroadcastStateNewSavePointsMain {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment streamExecutionEnvironment = FlinkEnvironmentFactory.simpleEnv(new EnvironmentConfig(null));

        DataStream<BroadcastStateDao> broadcastStateDaoDataStream = streamExecutionEnvironment.fromElements(
                new BroadcastStateDao(1, 0),
                new BroadcastStateDao(2, 1000),
                new BroadcastStateDao(3, 10000),
                new BroadcastStateDao(4, 100000),
                new BroadcastStateDao(5, 1000000)
        );

        StateBootstrapTransformation<BroadcastStateDao> broadcastTransformation = OperatorTransformation
                .bootstrapWith(broadcastStateDaoDataStream)
                .transform(new BroadcastStateNewSavePointsFunction());

        String newSavePointDirPathName = SavePointApiUtil.getNewSavePointWriterJobPath(CHECK_POINT_BASE_PATH + CUSTOM_BROADCAST_STATE_DIR);
        SavepointWriter
                .newSavepoint(new HashMapStateBackend(), MAX_PARALLELISM)
                .withOperator(BROADCAST_STATE_UID, broadcastTransformation)
                .write(newSavePointDirPathName);

        ModifyingPointDao modifyingPointDao = SavePointApiUtil.getModifyingPath(CHECK_POINT_BASE_PATH + CUSTOM_BROADCAST_STATE_DIR);
        SavepointWriter
                .fromExistingSavepoint(modifyingPointDao.getOldPath(), new HashMapStateBackend())
                .withOperator(BROADCAST_STATE_UID + MODIFYING_UID_PREFIX, broadcastTransformation)
                .write(modifyingPointDao.getNewPath());

        streamExecutionEnvironment.execute("BroadcastStateNewSavePointsMain");
    }

}
